"""Module to handle a structured grid in 1D

The Grid1D object is as simple as it can get. An equally spaced set of
nodes define the cell centers for a finite volume scheme. Each cell
has two faces at which a numerical flux must be computed.

The cell is represented by the `Cell1D` class and the faces by the
`Face1D` class. The `Grid1D` class stores the faces and cells for the
entire mesh.

"""
from grid import FaceType
from reconstruct1d import reconstruction_methods

import numpy

def _check(left, right):
    if left is None or right is None:
        msg = "Adjoining cells not defined"
        raise RuntimeError(msg)            

#############################################################################
# `Cell1D class`
#############################################################################
class Cell1D(object):
    r"""A 1D cell object.

    **Data Attributes**:

    grid : Grid1D
        The grid data structure owning the cell.

    left, right : Face1D
        The left and right faces for the cell.

    q, w : numpy.ndarray shape = (nvar, 1)
        The vector of conserved/primitive variables for the cell.

    index : int
        The cell index.

    dx : float
        Cell length.

    sigma : array (1, nvar)
        Approximation to the slope in this cell 
    
    """

    def __init__(self, index, grid, q, w):
        self.index = index
        self.grid = grid

        self.q = q
        self.w = w
        
        self.left = None
        self.right = None

        # set the cell length
        self.dx = grid.x[index+1] - grid.x[index]

        # set the default slope value
        self.sigma = numpy.zeros_like(q)

    def get_fluxes(self):
        """Get the fluxes for the faces bounding the cell.

        Returns : tuple (fl, fr)
            `fl` : flux for the left face.
            `fr` : flux for the right face.
        """
        _check(self.left, self.right)

        return (self.left.flux, self.right.flux)

    def get_sound_speed(self):
        """Get the sound speed for this cell.

        The cell index value is used to get the correct sound speed
        value.

        """
        return self.grid.a[self.index]        

#############################################################################
# `Face1D class`
#############################################################################    
class Face1D(object):
    """A 1D face object.

    Data Attributes:
    -----------------

    grid : Grid1D
        The grid data structure owning the cell.

    left, right : Cell1D
        left and right cells adjoining the face.

    index : int
        Index of the face.

    flux : numpy.ndarray
        The flux vector associated with the cell face.

    vertex_vel : float (0)
        An optional vertex velocity when the points move as in a
        Lagrangian simulation.

    """
    def __init__(self, index, grid, flux):
        self.index = index
        self.grid = grid
        
        self.left = None
        self.right = None

        self.flux = flux
        
        # set the vertex velocity
        self.vertex_vel = 0.0

    def get_primitive_variables(self, i = 0):
        r"""Return the primitive variables for the face.

        Parameters:
        ------------

        i : int
           Index to return the adjoining cells.

        For a face with index :math:`i + \frac{1}{2}`, the adjoining
        (right, left) cells are :math:`cell_{i+1}` and :math:`cell_i`.

        To implement high resolution methods, we need adjoining cells
        on the upwind side, where upwind is defined by the sign of an
        appropriate wave speed. Getting the two cell values on the
        upwind side is governed by the integer parameter :math:`i`.

        For :math:`i = 0`, the cells adjoining this face are
        returned. For :math:`i = 1`, the cells adjoining the next face
        is returned and simillarly, for :math:`i = -1`, the cell
        values for the previous face is returned.

        Returns: tuple (rho, u, p) `w` for each adjoining cell

        `wl` : primitive variables for the left cell
        `wr` : primitive variables for the right cell

        """
        # get the appropriate face
        face = self.grid.get_face(self.index + i)

        _check(face.left, face.right)

        wl = face.left.w; wr = face.right.w
        return (wl, wr)

    def get_conserved_variables(self, i = 0):
        r"""Return the conserved variables for the face.

        Parameters:
        ------------

        i : int
           Index to return the adjoining cells.

        For a face with index :math:`i + \frac{1}{2}`, the adjoining
        (right, left) cells are :math:`cell_{i+1}` and :math:`cell_i`.

        To implement high resolution methods, we need adjoining cells
        on the upwind side, where upwind is defined by the sign of an
        appropriate wave speed. Getting the two cell values on the
        upwind side is governed by the integer parameter :math:`i`.

        For :math:`i = 0`, the cells adjoining this face are
        returned. For :math:`i = 1`, the cells adjoining the next face
        is returned and simillarly, for :math:`i = -1`, the cell
        values for the previous face is returned.

        Returns : tuple (rho, M, E) `q` for each adjoining cell

        `ql` : conserved variables for the left cell.
        `qr` : conserved variables for the right cell.

        """
        face = self.grid.get_face(self.index + i)
        _check(face.left, face.right)

        ql = face.left.q; qr = face.right.q
        return (ql, qr)

    def get_sound_speeds(self, i=0):
        """Get the sound speed for ech adjoining face.

        Returns : tuple (cl, cr) for each adjoining face.

        """
        face = self.grid.get_face(self.index + i)
        _check(self.left, self.right)

        cl = face.left.get_sound_speed()
        cr = face.right.get_sound_speed()

        return (cl, cr)

    def compute_roe_average(self, left, right):
        """Compute the Roe's average of a quantity.

        Parameters:
        -----------

        left, right : float
            Property values on either side of the face.

        Notes:
        ------

        The Roe's average is defined as:

        ..math::

        f_{avg} = \frac{f_l} \sqrt{\rho_l} + f_r
        \sqrt{\rho_r}{sqrt{\rho_l} + \sqrt{\rho_r}}
        
        """
        # get the density on either side of the face
        ql, qr = self.get_conserved_variables()

        rhol, rhor = numpy.sqrt(ql[0]), numpy.sqrt(qr[0])
        
        val = left * rhol + right * rhor
        val /= (rhol + rhor)

        return val

    def update_position(self, dt):
        """Update the coordinate position of the face.

        Parameters:
        ------------

        dt : float
           The time step used for the update.

        Notes:
        ------

        The face is located at the grid's node locations (grid.x). The
        vertex velocity is defined at this point and at the end of a
        Lagrangian calculation, the cell (as bounded by the two faces)
        deforms based on the vertex velocity.

        """
        x = self.grid.x
        x[self.index] += self.vertex_vel * dt
        
#############################################################################
# `Grid1D class`
#############################################################################
class Grid1D(object):
    r"""General 1D structured mesh."""

    def __init__(self, reconstruction_method=0):
        # set a defalt reconstruction method
        self.reconstruction_method = reconstruction_method

    def set_reconstruction_function(self, reconstruction_method, **kwargs):
        self.reconstruction_method = reconstruction_method
        func = reconstruction_methods[reconstruction_method][1]

        self.reconstruction_function = func( self, **kwargs )

    def initialize(self, xlow, xhigh, dx, nvar, nb=2):
        r"""Initialize a 1D grid with equal spacing.

        **Parameters**:

        xlow : float
            Lower limit of the `x` domain

        xhigh : float
            Upper limit of the `x` domain

        dx : float
            Uniform spacing in the `x` direction.

        nvar : int
            Number of conserved variables.

        nb : int
            Number of boundary cells. Defaults to 2

        **Notes**

        Initializes the grid data structure. The first and last
        :math:`nb` cells are ghost cells used to imppose boundary
        conditions.

        The initialization also sets up the Cell1D and Face1D objects
        which comprise the cell. At the end of this function, the
        :class:`Grid1D` has a list of cells making up the mesh and a
        list of unique faces for the mesh.

        """

        # domain limits
        self.xlow = xlow
        self.xhigh = xhigh
        self.nb = nb

        # number of conserved variables
        self.nvar = nvar

        # number of cells
        ncells = int( float((xhigh - xlow))/float(dx) )
        
        # generate the grid points. The grid points determine the cell
        # faces in a cell-centered finite volume scheme.
        x = []
        for i in range(ncells + 1):
            x.append( xlow + i * dx )

        # add the boundary cells
        for i in range(nb):
            x.insert(0, x[0] - dx)
            x.append(x[-1] + dx)            

        # `x` is the spatial domain. The cell centers are midway
        # between consecutive points
        self.x = x = numpy.array(x)
        self.xc = 0.5 * (x[1:] + x[:-1])

        self.dx = dx

        # total number of nodes in the domain
        self.nnodes = len(self.x)
        self.ncells = len(self.xc)

        # initialize the vector of conserved variables
        self.q = numpy.zeros( shape=(nvar, self.ncells) )

        # initialize the vector of primitive variables 'w'
        self.w = self.q.copy()
        
        # create a list for the boundary cells and faces
        self.boundary_cells = {
                               'wall':{'cells':[], 'images':[]},
                               'periodic':{'cells':[], 'images':[]} }
        
        self.boundary_faces = {'left':[], 'right':[]}

        # setup the cells and faces
        self._setup_cells(nvar=nvar)    

    def get_face(self, index):
        """Get a face with the corresponding index"""
        assert 0 <= index < self.nfaces, "Invalid face index %d"%index
        return self.faces[index]

    def get_cell(self, index):
        """Get a cell with the corresponding index"""
        assert 0 <= index < self.ncells, "Invalid cell index %d"%index
        return self.cells[index]

    def _setup_cells(self, nvar):

        # initialize the flux vector
        self.nfaces = self.ncells + 1
        self._flux = numpy.zeros( shape=(nvar, self.nfaces) )

        # setup the cells
        self.cells = cells = []
        for i in range(self.ncells):
            self.cells.append( Cell1D(i, self, self.q[:,i], self.w[:,i]) )
            
        # store the boundary cells
        nb = self.nb
        
        bcells = []; images = []
        
        # wall boundary cells
        for i in  range(nb):
            # left
            bcells.append( cells[i] )
            images.append( cells[2*nb -1 -i] )
            
            # right
            bcells.append(cells[-nb + i])
            images.append(cells[-nb -1-i])
            
        self.boundary_cells["wall"]["cells"].extend( bcells)
        self.boundary_cells["wall"]["images"].extend( images )
        
        # periodic boundary cells
        bcells = []; images = []
        for i in range(nb):
            # left
            bcells.append( cells[i] )
            images.append( cells[-2*nb + i])
            
            #right
            bcells.append( cells[-nb + i] )
            images.append( cells[nb + i] )
            
        self.boundary_cells["periodic"]["cells"].extend( bcells )
        self.boundary_cells["periodic"]["images"].extend( images )
        
        # setup the faces
        self.faces = faces = []
        for i in range(self.nfaces):
            face = Face1D(i, self, self._flux[:,i])

            if i == 0:
                face.left = None

                cell_right = cells[i]
                cell_right.left = face
                
                face.right = cell_right
                
                # add this as a boundary face
                face.id = FaceType.Boundary
                self.boundary_faces['left'].append(face)

            elif i == self.ncells:
                cell_left = cells[i-1]
                cell_left.right = face
                
                face.left = cell_left
                face.right = None
                
                # add this as a boundary face
                face.id = FaceType.Boundary
                self.boundary_faces['right'].append(face)

            else:
                cell_left = cells[i-1]
                cell_right = cells[i]

                cell_left.right = face
                cell_right.left = face
                
                face.left = cell_left
                face.right = cell_right
                
                # this is an interior face
                face.id = FaceType.Interior

            faces.append(face)
